#include "Steganography.h"

/*
 *--------------------------------------------------------------------------------------
 *    Function:  hideMessage
 * Description:  Esconde uma mensagem
 *--------------------------------------------------------------------------------------
 */
int hideMessage(int argc, char *argv[])
{
    if(argc != 6)
    {
        cerr << "Usage mode: " << argv[0] << " hide <image filename> <message filename> <key filename> <output image filename>" << endl;
        return EXIT_FAILURE;
    }

    // Verifica a extensão da imagem de entrada
    string imfile(argv[2]);
    if(imfile.rfind(".bmp") == string::npos)
    {
        cerr << "Error: Input image should have the .bmp extension" << endl;
        return EXIT_FAILURE;
    }

    // Verifica a extensão da imagem de saída
    string outfile(argv[5]);
    if(outfile.rfind(".bmp") == string::npos)
    {
        cerr << "Error: Output image should have the .bmp extension" << endl;
        return EXIT_FAILURE;
    }

    // Lê a imagem de entrade aonde será a mensagem será escondida
    Image image = ImageIO::read(imfile);

    // Mensagem que será escondida
    ifstream msg(argv[3], ios_base::binary);
    if(!msg)
        throw runtime_error("Error: Could not open file " + string(argv[3]) + "\n");

    // Recupera a chave
    string keyfile(argv[4]);
    Key key(keyfile);

    // Esconde a mensagem
    Steganography::hide(image, msg, key);

    // Salva a imagem de saída
    ImageIO::write(image, outfile);

    return EXIT_SUCCESS;
}

/*
 *--------------------------------------------------------------------------------------
 *    Function:  extractMessage
 * Description:  Extrai uma mensagem
 *--------------------------------------------------------------------------------------
 */
int extractMessage(int argc, char *argv[])
{
    if(argc != 5)
    {
        cerr << "Usage mode: " << argv[0] << " extract <image filename> <key filename> <output message filename>" << endl;
        return EXIT_FAILURE;
    }

    // Verifica a extensão da imagem de entrada
    string imfile(argv[2]);
    if(imfile.rfind(".bmp") == string::npos)
    {
        cerr << "Error: Input image should have the .bmp extension" << endl;
        return EXIT_FAILURE;
    }

    // Lê a imagem de entrade aonde será a mensagem será escondida
    Image image = ImageIO::read(imfile);

    // Mensagem que será escondida
    ofstream msg(argv[4], ios_base::binary);
    if(!msg)
        throw runtime_error("Error: Could not open file " + string(argv[4]) + "\n");

    // Recupera a chave
    string keyfile(argv[3]);
    Key key(keyfile);

    // Recupera a mensagem
    Steganography::extract(image, msg, key);

    return EXIT_SUCCESS;
}

/*
 *--------------------------------------------------------------------------------------
 *    Function:  generateKey
 * Description:  Gera uma chave aleatória
 *--------------------------------------------------------------------------------------
 */
int generateKey(int argc, char *argv[])
{
    if(argc != 4)
    {
        cerr << "Usage mode: " << argv[0] << " gen-key <key filename> <quality=low|medium|high>" << endl;
        return EXIT_FAILURE;
    }

    string keyfile(argv[2]);
    int quality;

    // Verifica a qualidade passada
    if(!strcmp(argv[3], "low"))
        quality = Key::LOW_QUALITY;
    else if(!strcmp(argv[3], "medium"))
        quality = Key::MEDIUM_QUALITY;
    else if(!strcmp(argv[3], "high"))
        quality = Key::HIGH_QUALITY;
    else
    {
        cerr << "Error: Quality should be 'low', 'medium' or 'high'" << endl;
        return EXIT_FAILURE;
    }

    // Cria uma chave aleatória e escreve no arquivo de saída
    Key key(quality);
    key.write(keyfile);

    return EXIT_SUCCESS;
}

/*
 *--------------------------------------------------------------------------------------
 *    Function:  main
 * Description:  Função principal
 *--------------------------------------------------------------------------------------
 */
int main(int argc, char *argv[])
{
    // Quantidade inválida de argumentos
    if(argc < 2)
    {
        cerr << "Usage mode: " << argv[0] << " <hide|extract|gen-key> ..." << endl;
        return EXIT_FAILURE;
    }

    // Caso seja para esconder
    if(!strcmp(argv[1], "hide"))
        return hideMessage(argc, argv);

    // Caso seja para extrair
    else if(!strcmp(argv[1], "extract"))
        return extractMessage(argc, argv);

    // Caso seja para gerar uma chave
    else if(!strcmp(argv[1], "gen-key"))
        return generateKey(argc, argv);

    // Caso seja outra coisa
    else
    {
        cerr << "Usage mode: " << argv[0] << " <hide|extract|gen-key> ..." << endl;
        return EXIT_FAILURE;
    }
}

